#!/usr/bin/env python3

"""
EEG-ADC Analog-Digital Converter and Analyzer for scanned EEG recordings

Image properties:
100 DPI
490 x 100 px
H 2.5 cm
W 12.5 cm
"""

import sys
from os import path
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy
import pandas as pd
from scipy import interpolate
from scipy import ndimage
from scipy import misc
from scipy import signal
from scipy.fft import fft
from scipy.fft import fftfreq
from scipy.integrate import simpson
import skimage
from skimage import io
from skimage import morphology
from skimage import color
from skimage.filters import threshold_otsu
from skimage.filters import unsharp_mask

program_version = "1.0"
print(f"EEG-ADC analyzer v{program_version}")
print("=====================\n")

print("Imported modules:")
print(" Matplotlib " + matplotlib.__version__)
print("      NumPy " + np.__version__)
print("     pandas " + pd.__version__)
print("      SciPy " + scipy.__version__)
print("    skimage " + skimage.__version__ + "\n")

if (len(sys.argv) == 0):
    print("ERROR: No input file provided!")
    quit()

in_filename = sys.argv[1]

if (len(sys.argv) > 2):
    dpi = int(sys.argv[-1])
else:
    dpi = 100

# load input image file
# check if file can be opened
if not path.exists(in_filename):
    print(f"ERROR: File {in_filename} does not exist!")
    quit()

try:
    img = io.imread(in_filename, as_gray=True)
except FileNotFoundError:
    print(f"ERROR: File {in_filename} cannot be read!")
    quit
else:
    print(f"Reading input image: {in_filename}\n")

# denoise image with median filter
img_denoise = ndimage.median_filter(img, size=3)

# unsharp mask
img_sharp = unsharp_mask(img_denoise, radius=5, amount=2)

# edge detection
img_edges = ndimage.sobel(img_sharp)

# convert to b&w
thresh = threshold_otsu(img_edges)
img_bin = img_edges > thresh

# remove extreme data
img_bin[0, :] = True
img_bin[:, 0] = True
img_bin[-1, :] = True
img_bin[:, -1] = True

# interpolate
# TODO: input argument for user-defined px_uv and px_s
print("Interpolate and digitize..")
img_M = skimage.img_as_ubyte(img_bin)
px_cm = 0.03937008 * dpi * 10      # 1 dpi = 0.03937008 px/mm => px/mm => px/cm
px_uv = 200 / px_cm                # 1 cm = 200 uV => convert to pixels
px_s = 0.00025 * px_cm             # 2.5 cm = 1 s => convert to pixels
fs = dpi

dimx, dimy = img_M.shape

# fill empty columns
for ii in range(dimy):
    m = np.where(img_M[:, ii] == 0)[0]
    if len(m) == 0:
        img_M[:, ii] = np.zeros((dimx))

# interpolate column 0
m = np.where(img_M[:, 0] == 0)[0]
if len(m) == dimx:
    img_M[:, 0] = np.repeat(255, dimx)
    img_M[dimx // 2, 0] = 0

# interpolate other columns
for ii in range(dimy):
    m = np.where(img_M[:, ii] == 0)[0]
    if len(m) == dimx:
        img_M[:, ii] = img_M[:, ii - 1]

# convert to EEG signal
eeg_signal = np.zeros(dimy)
eeg_time = np.zeros(dimy)
for ii in range(dimy):
    eeg_signal[ii] = px_uv * (dimx // 2 - np.median(np.where(img_M[:, ii] == 0)[0]))
    eeg_time[ii] = ii * px_s

# interpolate EEG signal to remove artifacts
eeg_peaks = np.abs(np.diff(eeg_signal))
for ii in range(dimy-1):
    if eeg_peaks[ii] > np.mean(eeg_peaks) + 2 * np.std(eeg_peaks):
        eeg_signal[ii + 1] = np.mean((eeg_signal[ii], eeg_signal[ii + 2]))

# save bin image
bw_filename = in_filename.replace(".png", "-bw.png")
print(f"Saving PNG: {bw_filename}")
io.imsave(bw_filename, img_M)

# save signal to CSV file
df = pd.DataFrame({'time': eeg_time, 'signal': eeg_signal})
csv_filename = in_filename.replace(".png", "-signal.csv")
print(f"Saving CSV: {csv_filename}")
df.to_csv(csv_filename, index=False)

# HP filter at 0.5 Hz
# filter introduces delay, do not use filtered signal for plotting
print("Applying HP filter at 0.5 Hz..")
hp_fltr = signal.butter(10, 0.5, 'hp', fs=fs, output='sos')
eeg_filtered = signal.sosfilt(hp_fltr, eeg_signal)

# FFT
print("Calculating FFT..")
N = len(eeg_filtered)
# Hann tapering
h = signal.hann(N)
eeg_tapered = eeg_filtered * h
# remove DC from the signal
eeg_nodc = eeg_tapered - np.mean(eeg_tapered)
eeg_det = signal.detrend(eeg_nodc)
eeg_fft = fft(eeg_det)
eeg_fft_amp = 2 * np.abs(eeg_fft) / N
hz = np.linspace(0, fs // 2, (N // 2) + 1)

# Power Spectrum Density
print("Calculating PSD..")
freqs, psd = signal.welch(eeg_det, fs, nperseg=4*fs)

def band_power(psd, freqs, f1, f2):
    """
    Calculates absolute band power
    """
    # find intersecting values in frequency vector
    idx_delta = np.logical_and(freqs >= f1, freqs <= f2)
    # dx: frequency resolution
    p = simpson(psd[idx_delta], dx=freqs[1]-freqs[0])
    return p


# compute the absolute power by approximating the area under the curve
freq_res = freqs[1] - freqs[0]
total_power = simpson(psd, dx=freq_res)
abs_delta = band_power(psd, freqs, 0, 4)
abs_theta = band_power(psd, freqs, 4, 8)
abs_alpha = band_power(psd, freqs, 8, 13)
abs_alpha1 = band_power(psd, freqs, 8, 10)
abs_alpha2 = band_power(psd, freqs, 10, 13)
abs_beta = band_power(psd, freqs, 13, 31)
abs_beta1 = band_power(psd, freqs, 13, 18)
abs_beta2 = band_power(psd, freqs, 18, 31)
abs_gamma = band_power(psd, freqs, 31, hz[-1])
abs_gamma1 = band_power(psd, freqs, 31, 41)
abs_gamma2 = band_power(psd, freqs, 41, hz[-1])

rel_delta = abs_delta / total_power
rel_theta = abs_theta / total_power
rel_alpha = abs_alpha / total_power
rel_alpha1 = abs_alpha1 / total_power
rel_alpha2 = abs_alpha2 / total_power
rel_beta = abs_beta / total_power
rel_beta1 = abs_beta1 / total_power
rel_beta2 = abs_beta2 / total_power
rel_gamma = abs_gamma / total_power
rel_gamma1 = abs_gamma1 / total_power
rel_gamma2 = abs_gamma2 / total_power

# draw and save plots
pdf_filename = in_filename.replace(".png", "-signal.pdf")
print(f"Saving PDF: {pdf_filename}")
plt.figure(figsize=(12, 4))
plt.plot(eeg_time, eeg_signal, lw=1)
plt.xlim((0, eeg_time[-1]))
plt.ylim((-250, 250))
plt.title('EEG signal')
plt.xlabel('Time [s]')
plt.ylabel('Voltage [μV]')
plt.grid()
plt.savefig(pdf_filename, format="pdf", bbox_inches="tight")

# plot FFT
pdf_filename = in_filename.replace(".png", "-fft.pdf")
print(f"Saving PDF: {pdf_filename}")
plt.figure(figsize=(12, 4))
plt.plot(hz, eeg_fft_amp[0:len(hz)], lw=1)
plt.xlabel('Frequency (Hz)')
plt.ylabel('2 × |FFT|')
plt.xlim([0, freqs.max()])
plt.ylim([0, np.max(eeg_fft_amp[0:len(hz)]) * 1.1])
plt.text(2,
         np.max(eeg_fft_amp[0:len(hz)]),
         'delta',
         horizontalalignment='center')
plt.axvline(x=4,
            color='gray',
            linestyle='--')
plt.text(6,
         np.max(eeg_fft_amp[0:len(hz)]),
         'theta',
         horizontalalignment='center')
plt.axvline(x=8,
            color='gray',
            linestyle='--')
plt.text(10.5,
         np.max(eeg_fft_amp[0:len(hz)]),
         'alpha',
         horizontalalignment='center')
plt.axvline(x=13,
            color='gray',
            linestyle='--')
plt.text(20.5,
         np.max(eeg_fft_amp[0:len(hz)]),
         'beta',
         horizontalalignment='center')
plt.axvline(x=30,
            color='gray',
            linestyle='--')
plt.text(40,
         np.max(eeg_fft_amp[0:len(hz)]),
         'gamma',
         horizontalalignment='center')
plt.title("FFT")
plt.grid()
plt.savefig(pdf_filename, format="pdf", bbox_inches="tight")

# plot the power spectrum
pdf_filename = in_filename.replace(".png", "-psd.pdf")
print(f"Saving PDF: {pdf_filename}")
plt.figure(figsize=(12, 4))
plt.plot(freqs, psd, lw=1)
plt.xlabel('Frequency (Hz)')
plt.ylabel('Power spectral density (μV^2 / Hz)')
plt.xlim([0, freqs.max()])
plt.ylim([0, psd.max() * 1.1])
plt.text(2, psd.max() + 0.5,
         'delta', horizontalalignment='center')
plt.axvline(x=4, color='gray', linestyle='--')
plt.text(6, psd.max() + 0.5,
         'theta', horizontalalignment='center')
plt.axvline(x=8, color='gray', linestyle='--')
plt.text(10.5, psd.max() + 0.5,
         'alpha', horizontalalignment='center')
plt.axvline(x=13, color='gray', linestyle='--')
plt.text(20.5, psd.max() + 0.5,
         'beta', horizontalalignment='center')
plt.axvline(x=30, color='gray', linestyle='--')
plt.text(40, psd.max() + 0.5,
         'gamma', horizontalalignment='center')
plt.title("Welch's periodogram")
plt.grid()
plt.savefig(pdf_filename, format="pdf", bbox_inches="tight")

# plot the absolute powers
pdf_filename = in_filename.replace(".png", "-bands.pdf")
print(f"Saving PDF: {pdf_filename}")
abs = [total_power, abs_delta, abs_theta, abs_alpha,
abs_alpha1, abs_alpha2, abs_beta, abs_beta1,
abs_beta2, abs_gamma, abs_gamma1, abs_gamma2]
bands = ["total power", "delta", "theta", "alpha", "alpha1", "alpha2", "beta", "beta1", "beta2", "gamma", "gamma1", "gamma2"]
plt.figure(figsize=(12, 4))
plt.bar(bands, abs, lw=1)
plt.ylabel('Absolute power [μV^2/Hz]')
plt.title('Absolute power across EEG bands')
plt.savefig(pdf_filename, format="pdf", bbox_inches="tight")

# save band data
bands = ['total_power', 'delta', 'theta', 'alpha', 'alpha1', 'alpha2', 'beta', 'beta1', 'beta2', 'gamma', 'gamma1', 'gamma2']
abs_power = [total_power, abs_delta, abs_theta, abs_alpha, abs_alpha1, abs_alpha2, abs_beta, abs_beta1, abs_beta2, abs_gamma, abs_gamma1, abs_gamma2]
df = pd.DataFrame({'band': bands, 'abs_power': abs_power})
csv_filename = in_filename.replace(".png", "-bands.csv")
print(f"Saving CSV: {csv_filename}")
df.to_csv(csv_filename, index=False)

# write report
txt_filename = in_filename.replace(".png", "-report.txt")
print(f"Saving TXT: {txt_filename}\n")
with open(txt_filename, "w") as f:
    # header
    print("Signal data", file=f)
    print("===========", file=f)
    print(f"     Input signal: {in_filename}", file=f)
    print(f"Signal length [s]: {round(eeg_time[-1], 2)}\n", file=f)

    # amplitude
    print("Signal amplitude", file=f)
    print("================", file=f)
    print(f"   min amplitude [μV]: {round(np.min(eeg_filtered), 2)}", file=f)
    print(f"   max amplitude [μV]: {round(np.max(eeg_filtered), 2)}", file=f)
    print(f"  mean amplitude [μV]: {round(np.mean(eeg_filtered), 2)}\n", file=f)
    print(f"median amplitude [μV]: {round(np.median(eeg_filtered), 2)}\n", file=f)

    # power in EEG ranges
    print("Signal power [μV^2]", file=f)
    print("===================", file=f)
    print(f"                      total power: {round(total_power, 2)}", file=f)
    print(f"   δ band (0-3 Hz) absolute power: {round(abs_delta, 2)}", file=f)
    print(f"            δ band relative power: {round(rel_delta, 2)}", file=f)
    print(f"   θ band (4-7 Hz) absolute power: {round(abs_theta, 2)}", file=f)
    print(f"            θ band relative power: {round(rel_theta, 2)}", file=f)
    print(f"  α band (8-12 Hz) absolute power: {round(abs_alpha, 2)}", file=f)
    print(f"            α band relative power: {round(rel_alpha, 2)}", file=f)
    print(f"  α1 band (8-9 Hz) absolute power: {round(abs_alpha1, 2)}", file=f)
    print(f"           α1 band relative power: {round(rel_alpha1, 2)}", file=f)
    print(f"α2 band (10-12 Hz) absolute power: {round(abs_alpha2, 2)}", file=f)
    print(f"           α2 band relative power: {round(rel_alpha2, 2)}", file=f)
    print(f" β band (13-30 Hz) absolute power: {round(abs_beta, 2)}", file=f)
    print(f"            β band relative power: {round(rel_beta, 2)}", file=f)
    print(f"β1 band (13-17 Hz) absolute power: {round(abs_beta1, 2)}", file=f)
    print(f"           β1 band relative power: {round(rel_beta1, 2)}", file=f)
    print(f"β2 band (18-30 Hz) absolute power: {round(abs_beta2, 2)}", file=f)
    print(f"           β2 band relative power: {round(rel_beta2, 2)}", file=f)
    print(f" γ band (31-{round(hz[-1])} Hz) absolute power: {round(abs_gamma, 2)}", file=f)
    print(f"            γ band relative power: {round(rel_gamma, 2)}", file=f)
    print(f"γ1 band (31-40 Hz) absolute power: {round(abs_gamma1, 2)}", file=f)
    print(f"           γ1 band relative power: {round(rel_gamma1, 2)}", file=f)
    print(f"γ2 band (42-{round(hz[-1])} Hz) absolute power: {round(abs_gamma2, 2)}", file=f)
    print(f"           γ2 band relative power: {round(rel_gamma2, 2)}", file=f)

print("Analysis completed.")
